package org.gwu.seuss.communication.crypto;

import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.UnsupportedEncodingException;
import java.security.GeneralSecurityException;
import java.security.InvalidKeyException;
import java.security.KeyStore;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.security.SecureRandom;

import javax.crypto.Cipher;
import javax.crypto.KeyGenerator;
import javax.crypto.Mac;
import javax.crypto.SecretKey;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;

import org.apache.commons.codec.binary.Base64;
import org.apache.commons.io.FileUtils;
import org.apache.log4j.Logger;
import org.gwu.seuss.communication.exception.SecretKeyException;
import org.springframework.util.Assert;

public class CryptoUtil {
	static Logger log = Logger.getLogger(CryptoUtil.class);
	
	private static final String ENC_ALGORITHM = "AES";
	private static final String TRANS_ALGORITHM = "AES/CBC/PKCS5PADDING";
	private static final String HMAC_ALGORITHM = "HmacSHA256";
	private static final String HASH_ALGORITHM = "SHA-512";
	private static final String PRNG_ALGORITHM = "SHA1PRNG";
	
	private static final String SEUSS_DIR = "seuss";
	private static final String SECRET_KEY_ALIAS = "secretKeyAlias";
	private static final String KEYSTORE_EXT= "JCEKS";
	private static final String KEYSTORE_NAME= "keystore.jceks";
	
	private final static int ITERATION_NUMBER = 1000;
	
	private SecretKey secretKey = null;
	
	private static volatile CryptoUtil instance = null;
	
	private CryptoUtil() {		
	}
	
	public static CryptoUtil getInstance() {
		if (instance == null) {
			synchronized (CryptoUtil.class) {
				instance = new CryptoUtil();
			}
		}
		return instance;
	}

	/**
	 * 
	 * @param file
	 * @return
	 * @throws Exception
	 */
	public CryptoVal encrypt(SecretKey key, byte[] file) throws Exception {
		try {
			Cipher encrypt = Cipher.getInstance(TRANS_ALGORITHM);
			encrypt.init(Cipher.ENCRYPT_MODE, key);
			
			byte[] iv = encrypt.getIV();
			
			byte[] encFile = encrypt.doFinal(file);
			
			CryptoVal val = new CryptoVal();
			val.setEncFile(encFile);
			val.setIv(iv);;
			val.setKey(key);
			
			return val;
			
		} catch (Exception e) {
			log.error("Failed to encrypt file.", e);
			throw e;
		}
	}
	
	public SecretKey genUserSessionKey() throws NoSuchAlgorithmException {
		// Generate the key
		try {
			KeyGenerator gen = KeyGenerator.getInstance(ENC_ALGORITHM);
			gen.init(256);
			// k is secret and must be protected
			return gen.generateKey();
		} catch (NoSuchAlgorithmException e) {
			log.error("Failed to generate user session key.", e);
			throw e;
		}

	}
	
	public SecretKey getSecretKey(byte[] bytes) {
		return new SecretKeySpec(bytes, 0, bytes.length, "AES");
	}
	
	
	/**
	 * Decrypt a file
	 * 
	 * @param iv
	 * @param key
	 * @param encFile
	 * @return
	 * @throws Exception
	 */
	public byte[] decrypt(byte[] iv, SecretKey key, byte[] encFile) throws Exception  {
		try {
			Cipher decrypt = Cipher.getInstance(TRANS_ALGORITHM);
			decrypt.init(Cipher.DECRYPT_MODE, key, new IvParameterSpec(iv));
			return decrypt.doFinal(encFile);
		} catch (Exception e) {
			log.error("Failed to decrypt file.", e);
			throw e;
		}
	}
	
	public String getPasswordHash(String password, byte[] salt) throws GeneralSecurityException {	
		try{			
			MessageDigest digest = MessageDigest.getInstance(HASH_ALGORITHM);
			digest.reset();
			digest.update(salt);
	
			byte[] input = digest.digest(password.getBytes("UTF-8"));

			for (int i = 0; i < ITERATION_NUMBER; i++) {
				digest.reset();
				input = digest.digest(input);
			}
	
			return Base64.encodeBase64String(input);
		} catch (NoSuchAlgorithmException | UnsupportedEncodingException e) {
			log.error("Failed to hash password.", e);
			throw new GeneralSecurityException();
		}	
	}
	
	/**
	 * 
	 * @param secret
	 * @param fileBytes
	 * @return
	 * @throws GeneralSecurityException
	 */
	public String generateHmac(SecretKey secret, byte[] fileBytes) throws GeneralSecurityException {
		try {
			Mac sha256Hmac = Mac.getInstance(HMAC_ALGORITHM);
			SecretKeySpec secret_key = new SecretKeySpec(secret.getEncoded(), HMAC_ALGORITHM);
			sha256Hmac.init(secret_key);			
			byte[] hmacBytes = sha256Hmac.doFinal(fileBytes);
			return Base64.encodeBase64String(hmacBytes);
			
		} catch (NoSuchAlgorithmException | InvalidKeyException e) {
			log.error("Failed to generate MAC for file.", e);
			throw e;
		}
	}
	
	/**
	 * Generate a 64-bit random salt value.
	 * @return
	 * @throws GeneralSecurityException
	 */
	public byte[] generateSalt() throws GeneralSecurityException {
		SecureRandom random = SecureRandom.getInstance(PRNG_ALGORITHM);
		
		// Salt generation 64 bits long
		byte[] bSalt = new byte[8];
		random.nextBytes(bSalt);
		return bSalt;
	}
	
	/**
	 * 
	 * @param secretKey
	 * @param password
	 * @throws Exception
	 */
	public void saveKeyStore(SecretKey secretKey, String password) throws Exception {
		
		Assert.notNull(secretKey, "User secret key is empty.");
		Assert.hasText(password, "Password is not provided");

		KeyStore.ProtectionParameter protParam = new KeyStore.PasswordProtection(password.toCharArray());
		KeyStore.SecretKeyEntry skEntry = new KeyStore.SecretKeyEntry(secretKey);

		KeyStore keyStore = KeyStore.getInstance(KEYSTORE_EXT);
		keyStore.load(null, password.toCharArray());
		keyStore.setEntry(SECRET_KEY_ALIAS, skEntry, protParam);

		FileOutputStream fos = null;
		try {
			fos = FileUtils.openOutputStream(new File(SEUSS_DIR + File.separatorChar + KEYSTORE_NAME));
			keyStore.store(fos, password.toCharArray());
		} finally {
			if (fos != null) {
				fos.close();
			}
		}
	}
	
	/**
	 * 
	 * @param directory
	 * @param newPassword
	 * @param oldPassword
	 * @throws Exception 
	 */
	public void replaceKeyStore(String newPassword, String oldPassword) throws Exception {
		
		Assert.hasText(newPassword, "New password is not provided");
		Assert.hasText(oldPassword, "Old password is not provided");
	
		// Need to load the secret key with old password
		SecretKey loadedKey = getSecretKeyFromKeyStore(oldPassword);
		
		// Override the old key store with a new key store with new password
		saveKeyStore(loadedKey, newPassword);
	}
	
	/**
	 * 
	 * @param password
	 * @return
	 * @throws GeneralSecurityException
	 * @throws IOException
	 */
	private SecretKey getSecretKeyFromKeyStore(String password) throws GeneralSecurityException, IOException {
		KeyStore keyStore = KeyStore.getInstance(KEYSTORE_EXT);
		
		java.io.FileInputStream fis = null;
	    try {
	        fis = FileUtils.openInputStream(new File(SEUSS_DIR + File.separatorChar + KEYSTORE_NAME));
	        keyStore.load(fis, password.toCharArray());
	    } finally {
	        if (fis != null) {
	            fis.close();
	        }
	    }
		
		KeyStore.ProtectionParameter protParam = new KeyStore.PasswordProtection(password.toCharArray());
		KeyStore.SecretKeyEntry pkEntry = (KeyStore.SecretKeyEntry) keyStore.getEntry(SECRET_KEY_ALIAS, protParam);
		
		return pkEntry.getSecretKey();
	}
	
	/**
	 * 
	 * @param password
	 * @throws Exception
	 */
	public void createKeyStore(String password) throws Exception {
		Assert.hasText(password,  "Password is not provided");
		
		File keystore = new File(SEUSS_DIR, KEYSTORE_NAME);
		if(keystore.exists()) {
			secretKey = getSecretKeyFromKeyStore(password);
		} else {
			secretKey = genUserSessionKey();
			saveKeyStore(secretKey, password);
		}
	}

	/**
	 * 
	 * @return
	 * @throws SecretKeyException 
	 */
	public SecretKey getSecretKey() throws SecretKeyException {
		if (secretKey == null) {
			throw new SecretKeyException("Failed to load secret key.");
		}
		return secretKey;
	}
	
}